import paddle
import matplotlib.pyplot as plt
from visualdl import LogWriter
import platform
import os
if (platform.system() == 'Linux'):
    os.chdir("gan_paddle/")

# from Module import *
# from Util import *

from gan_paddle.Module import *
from gan_paddle.Util import *

import cv2

from parl.utils import logger

# 噪声维度
LR = 2e-4
Z_DIM = 100
BATCH_SIZE = 256
MODEL_PATH = './module/'
EPOCH = 50


# 读取真实图片的数据集，这里去除了数据集中的label数据，因为label在这里使用不上，这里不考虑标签分类问题。
def mnist_reader(reader):
    def r():
        for img, label in reader():
            yield img.reshape(1, 28, 28)

    return r


# 噪声生成，通过由噪声来生成假的图片数据输入。
def z_reader():
    while True:
        yield np.random.normal(0.0, 1.0, (Z_DIM, 1, 1)).astype('float32')


def test():
    # 生成真实图片reader
    mnist_generator = paddle.batch(
        paddle.reader.shuffle(mnist_reader(paddle.dataset.mnist.train()), 30000), batch_size=BATCH_SIZE)
    # 生成假图片的reader
    z_generator = paddle.batch(z_reader, batch_size=BATCH_SIZE)
    pics_tmp = next(mnist_generator())
    print('一个batch图片数据的形状：batch_size =', len(pics_tmp), ', data_shape =', pics_tmp[0].shape)
    plt.imshow(pics_tmp[0][0])
    plt.show()
    z_tmp = next(z_generator())
    print('一个batch噪声z的形状：batch_size =', len(z_tmp), ', data_shape =', z_tmp[0].shape)

    # 测试生成网络G和判别网络D
    # with fluid.dygraph.guard():
    #     g_tmp = G('G')
    #     tmp_g = g_tmp(fluid.dygraph.to_variable(np.array(z_tmp))).numpy()
    #     print('生成器G生成图片数据的形状：', tmp_g.shape)
    #     plt.imshow(tmp_g[0][0])
    #     plt.show()
    #
    #     d_tmp = D('D')
    #     tmp_d = d_tmp(fluid.dygraph.to_variable(tmp_g)).numpy()
    #     print('判别器D判别生成的图片的概率数据形状：', tmp_d.shape)
    # show_image_grid(tmp_g, BATCH_SIZE)
    # plt.figure(figsize=(12,BATCH_SIZE/32),dpi=80)
    # plt.imshow(concatenate_img(tmp_g))


class TrainAgent(object):
    def __init__(self, epoch_num=10, batch_size=128, use_gpu=True, load_model=False):
        self.epoch_num = epoch_num
        self.batch_size = batch_size
        self.load_model = load_model

        # 生成真实图片reader
        self.mnist_generator = paddle.batch(
            paddle.reader.shuffle(mnist_reader(paddle.dataset.mnist.train()), 30000), batch_size=BATCH_SIZE)
        # 生成假图片的reader
        self.z_generator = paddle.batch(z_reader, batch_size=BATCH_SIZE)
        self.place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()

        with fluid.dygraph.guard(self.place):
            # 模型存储路径
            self.model_path = MODEL_PATH

            self.d = D('D')
            self.d.train()
            self.g = G('G')
            self.g.train()

            # 创建优化方法
            self.real_d_optimizer = fluid.optimizer.AdamOptimizer(learning_rate=LR, parameter_list=self.d.parameters())
            self.fake_d_optimizer = fluid.optimizer.AdamOptimizer(learning_rate=LR, parameter_list=self.d.parameters())
            self.g_optimizer = fluid.optimizer.AdamOptimizer(learning_rate=LR, parameter_list=self.g.parameters())

            # 读取上次保存的模型
            if self.load_model == True:
                self.loadModel()

    def loadModel(self):
        logger.info('loadModel')
        g_para, g_opt = fluid.load_dygraph(self.model_path + 'g')
        d_para, d_r_opt = fluid.load_dygraph(self.model_path + 'd_o_r')
        # 上面判别器的参数已经读取到d_para了,此处无需再次读取
        _, d_f_opt = fluid.load_dygraph(self.model_path + 'd_o_f')

        self.g.load_dict(g_para)
        self.g_optimizer.set_dict(g_opt)
        self.d.load_dict(d_para)
        self.real_d_optimizer.set_dict(d_r_opt)
        self.fake_d_optimizer.set_dict(d_f_opt)

    def saveModel(self):
        model_path = self.model_path
        fluid.save_dygraph(self.g.state_dict(), model_path + 'g')
        fluid.save_dygraph(self.g_optimizer.state_dict(), model_path + 'g')

        fluid.save_dygraph(self.d.state_dict(), model_path + 'd_o_r')
        fluid.save_dygraph(self.real_d_optimizer.state_dict(), model_path + 'd_o_r')

        fluid.save_dygraph(self.d.state_dict(), model_path + 'd_o_f')
        fluid.save_dygraph(self.fake_d_optimizer.state_dict(), model_path + 'd_o_f')

    def train(self):
        with fluid.dygraph.guard(self.place):
            # 定义日志写入(先清空日志文件夹)
            # if load_model == False:
            #     !rm - rf / home / aistudio / log /
            # 创建 LogWriter 对象
            # log_writer = LogWriter("./log", sync_cycle=10)
            # real_loss_wrt = log_writer.scalar("d_loss")
            # fake_loss_wrt = log_writer.scalar("d_fake_loss")
            # g_loss_wrt = log_writer.scalar("g_loss")
            #image_wrt = log_writer.image("imgs", 1)

            iteration_num = 0
            logger.info('start')
            for epoch in range(self.epoch_num):
                for i, real_image in enumerate(self.mnist_generator()):
                    # 丢弃不满整个batch_size的数据
                    if len(real_image) != self.batch_size:
                        continue

                    iteration_num += 1

                    '''
                    判别器d通过最小化输入真实图片时判别器d的输出与真值标签ones的交叉熵损失，来优化判别器的参数，
                    以增加判别器d识别真实图片real_image为真值标签ones的概率。
                    '''
                    # 将MNIST数据集里的图片读入real_image，将真值标签ones用数字1初始化
                    real_image = fluid.dygraph.to_variable(np.array(real_image))
                    ones = fluid.dygraph.to_variable(np.ones([len(real_image), 1]).astype('float32'))
                    # 计算判别器d判断真实图片的概率
                    p_real = self.d(real_image)
                    # 计算判别真图片为真的损失
                    real_cost = fluid.layers.sigmoid_cross_entropy_with_logits(p_real, ones)
                    real_avg_cost = fluid.layers.mean(real_cost)
                    # 反向传播更新判别器d的参数
                    real_avg_cost.backward()
                    self.real_d_optimizer.minimize(real_avg_cost)
                    self.d.clear_gradients()

                    '''
                    判别器d通过最小化输入生成器g生成的假图片g(z)时判别器的输出，与假值标签zeros的交叉熵损失，
                    来优化判别器d的参数，以增加判别器d识别生成器g生成的假图片g(z)为假值标签zeros的概率。
                    '''
                    # 创建高斯分布的噪声z，将假值标签zeros初始化为0
                    z = next(self.z_generator())
                    z = fluid.dygraph.to_variable(np.array(z))
                    zeros = fluid.dygraph.to_variable(np.zeros([len(real_image), 1]).astype('float32'))
                    # 判别器d判断生成器g生成的假图片的概率
                    p_fake = self.d(self.g(z))
                    # 计算判别生成器g生成的假图片为假的损失
                    fake_cost = fluid.layers.sigmoid_cross_entropy_with_logits(p_fake, zeros)
                    fake_avg_cost = fluid.layers.mean(fake_cost)
                    # 反向传播更新判别器d的参数
                    fake_avg_cost.backward()
                    self.fake_d_optimizer.minimize(fake_avg_cost)
                    self.d.clear_gradients()

                    '''
                    生成器g通过最小化判别器d判别生成器生成的假图片g(z)为真的概率d(fake)，与真值标签ones的交叉熵损失，
                    来优化生成器g的参数，以增加生成器g使判别器d判别其生成的假图片g(z)为真值标签ones的概率。
                    '''
                    # 生成器用输入的高斯噪声z生成假图片
                    fake = self.g(z)
                    # 计算判别器d判断生成器g生成的假图片的概率
                    p_confused = self.d(fake)
                    # 使用判别器d判断生成器g生成的假图片的概率与真值ones的交叉熵计算损失
                    g_cost = fluid.layers.sigmoid_cross_entropy_with_logits(p_confused, ones)
                    g_avg_cost = fluid.layers.mean(g_cost)
                    # 反向传播更新生成器g的参数
                    g_avg_cost.backward()
                    self.g_optimizer.minimize(g_avg_cost)
                    self.g.clear_gradients()

                    # 打印输出

                        # show_image_grid(fake.numpy(), self.batch_size, epoch)

                    # real_loss_wrt.add_record(iteration_num, real_avg_cost.numpy()[0])
                    # fake_loss_wrt.add_record(iteration_num, fake_avg_cost.numpy()[0])
                    # g_loss_wrt.add_record(iteration_num, g_avg_cost.numpy()[0])
                    # 写日志
                    if iteration_num % 1000 == 0:
                        cv2.imwrite("./log/" + str(iteration_num) + '.jpg', concatenate_img(fake.numpy(), BATCH_SIZE))

                    # image_wrt.start_sampling()
                    # # 获取idx
                    # idx = image_wrt.is_sample_taken()
                    # # 如果 idx != -1，采样，否则跳过
                    # if idx != -1:
                    #     # 获取图片数据
                    #     image_data = concatenate_img(fake.numpy())
                    #     # 使用 set_sample() 函数添加数据
                    #     # flatten() 用于把 ndarray 由矩阵变为行向量
                    #     image_wrt.set_sample(idx, image_data.shape, image_data.flatten())
                    #     # 如果完成了当前轮的采样，则调用finish_sample()
                    #     image_wrt.finish_sampling()

                logger.info('epoch = {}, real_d_loss = {:.2f}, fake_d_loss = {:.2f}, g_loss = {:.2f}'.format(
                    epoch, real_avg_cost.numpy()[0], fake_avg_cost.numpy()[0], g_avg_cost.numpy()[0]))
            # 存储模型
            self.saveModel()

    def test(self):
        with fluid.dygraph.guard(self.place):
            z = next(self.z_generator())
            z = fluid.dygraph.to_variable(np.array(z))
            fake = self.g(z)
            cv2.imwrite('./log/test.jpg', concatenate_img(fake.numpy(), BATCH_SIZE))

if __name__ == '__main__':
    logger.auto_set_dir('d')

    agent = TrainAgent(epoch_num=EPOCH, batch_size=BATCH_SIZE, use_gpu=True, load_model=True)
    agent.test()

    # agent = TrainAgent(epoch_num=20, batch_size=BATCH_SIZE, use_gpu=True)
    # agent.train()
